from sklearn import neighbors
import os
import os.path
import pickle
from PIL import Image, ImageDraw, ImageFont
import face_recognition
from face_recognition.face_recognition_cli import image_files_in_folder




ALLOWED_EXTENSIONS = {'bmp', 'png', 'jpg', 'jpeg'}

num1=0

# 对指定的训练图片文件夹进行训练
def train(train_dir, model_save_path=None, n_neighbors=None, knn_algo='ball_tree', verbose=False):

    X = []
    y = []

    # 循环获取训练集图片
    for class_dir in os.listdir(train_dir):
        if not os.path.isdir(os.path.join(train_dir, class_dir)):
            continue
        # 遍历当前任务的每一张图片
        # image_files_in_folder，这个地方是获取文件夹下的所有图片文件，可以修改其中的图片类型
        # 默认是jpg|jpeg|png，后来追加了bmp
        print('training picture of {}'.format(class_dir))
        flag=0
        for img_path in image_files_in_folder(os.path.join(train_dir, class_dir)):
            # 加载图片文件，其实是numpy数组
            image = face_recognition.load_image_file(img_path)
            # 获取人脸检测框
            face_bounding_boxes = face_recognition.face_locations(image)
            # 多个人物或者0个人物不处理
            if len(face_bounding_boxes) != 1:
                # If there are no people (or too many people) in a training image, skip the image.
                if verbose:
                    print("Image {} not suitable for training: {}".format(img_path, "Didn't find a face" if len(
                        face_bounding_boxes) < 1 else "Found more than one face"))
            else:

                # Add face encoding for current image to the training set
                X.append(face_recognition.face_encodings(image, known_face_locations=face_bounding_boxes)[0])
                y.append(class_dir)
            flag+=1
            print(flag)
    # 设置KNN分类器的近邻数
    # Determine how many neighbors to use for weighting in the KNN classifier
    if n_neighbors is None:
        # n_neighbors = int(round(math.sqrt(len(X))))
        n_neighbors = 3
        if verbose:
            print("Chose n_neighbors automatically:", n_neighbors)

    # 创建KNN分类器，并进行训练
    knn_clf = neighbors.KNeighborsClassifier(n_neighbors=n_neighbors, algorithm=knn_algo, weights='distance')
    #n_neighbors 默认情况下kneighbors查询使用的邻居数。就是k-NN的k的值，选取最近的k个点
    #algorithm 快速k近邻搜索算法，默认参数为auto，可以理解为算法自己决定合适的搜索算法。
    #weights 预测的权函数，概率值。distance是距离越近权重越大
    knn_clf.fit(X, y)

    # 保存KNN训练结果
    if model_save_path is not None:
        with open(model_save_path, 'wb') as f:
            pickle.dump(knn_clf, f)

    return knn_clf


# 对指定的预测图片进行预测
def predict(X_img_path, knn_clf=None, model_path=None, distance_threshold=0.5):

    # 校验当前文件类型
    if not os.path.isfile(X_img_path) or os.path.splitext(X_img_path)[1][1:] not in ALLOWED_EXTENSIONS:
        raise Exception("Invalid image path: {}".format(X_img_path))

    # 校验当前模型文件和knn模型，两个不能全空
    if knn_clf is None and model_path is None:
        raise Exception("Must supply knn classifier either thourgh knn_clf or model_path")

    # 加载训练好的KNN模型
    if knn_clf is None:
        with open(model_path, 'rb') as f:
            knn_clf = pickle.load(f)

    # 加载图片，获取人脸检测框
    X_img = face_recognition.load_image_file(X_img_path)
    X_face_locations = face_recognition.face_locations(X_img)
    # print('predict {}'.format(X_img_path))
    # 如果未找到人脸，返回[]
    if len(X_face_locations) == 0:
        return []

    # 对测试图片进行编码转换，转换为numpy数组
    faces_encodings = face_recognition.face_encodings(X_img, known_face_locations=X_face_locations)

    # 通过KNN模型找到最佳匹配的人脸
    closest_distances = knn_clf.kneighbors(faces_encodings, n_neighbors=3)
    # # 返回值indices：第0列元素为参考点的索引，后面是(n_neighbors - 1)个与之最近的点的索引
    # # 返回值distances：第0列元素为与自身的距离(为0)，后面是(n_neighbors - 1)个与之最近的点与参考点的距离
    # closest_distances= [[0.34997745 0.3750366  0.37819395]]
    # closest_distances= [[ 5 12 11]]
    # for i in closest_distances:
    #     print('closest_distances=',i)
    are_matches = []
    # are_matches = [closest_distances[0][i][0] <= distance_threshold for i in range(len(X_face_locations))]
    for i in range(len(X_face_locations)):
        are_matches.append(closest_distances[0][i][0] <= distance_threshold)
        # print('predict value=', closest_distances[0][i][0])
    # print('knn_clf.predict(faces_encodings)=',knn_clf.predict(faces_encodings))
    # 预测分类
    return [(pred, loc) if rec else ("unknown", loc) for pred, loc, rec in
            zip(knn_clf.predict(faces_encodings), X_face_locations, are_matches)]


# 在图片上输出预测标签
def show_prediction_labels_on_image(img_path, predictions,img_path_test,img_path_res):

    # 打开图片，获取句柄
    pil_image = Image.open(img_path).convert("RGB")
    #创建一个写对象
    draw = ImageDraw.Draw(pil_image)

    # 对预测结果进行遍历
    for name, (top, right, bottom, left) in predictions:
        # 在人脸周边画矩形框
        draw.rectangle(((left, top), (right, bottom)), outline=(0, 0, 255))

        # There's a bug in Pillow where it blows up with non-UTF-8 text
        # when using the default bitmap font

        # 字体  字体*.ttc的存放路径一般是： /usr/share/fonts/opentype/noto/ 查找指令locate *.ttc
        font = ImageFont.truetype('simsun.ttc', 40)
        # 在人脸地图输出标签
        text_width, text_height = draw.textsize(name)
        draw.rectangle(((left, bottom - text_height - 10), (right, bottom)), fill=(0, 0, 255), outline=(0, 0, 255))
        #draw = ImageDraw.Draw(img_PIL)
        draw.text((left + 6, bottom - text_height - 5),name,font=font,fill=(255, 255, 255, 255))
    # Remove the drawing library from memory as per the Pillow docs
    print(draw)
    print(img_path_res)
    img_path_res='D:/pythonProject/djangoProject2'+img_path_res
    #del draw
    # Display the resulting image
    pil_image.save(img_path_res)



#本函数是多人人脸对比的函数，应该将识别出来的图存进img下的文件夹中，并且返回识别出来的人数

def face_detect1(img_path_test,img_path_res):
    #return_name=[]
    # 指定训练集路径和预测集路径
    #train_dir = '../static/img/facedetect/train'
    # train_dir = r'D:\pythonProject\djangoProject1\static\img\facedetect\train'
    # #test_dir = '../static/img/facedetect/test'
    test_dir = r'D:\pythonProject\djangoProject2\static\img\facedetect\test'
    # # 第一步，通过KNN分类器进行训练，并存储模型文件
    # print("Training KNN classifier...")
    # classifier = train(train_dir, model_save_path="trained_knn_model.clf", n_neighbors=3)
    # print("Training complete!")
    # 第二步，使用训练分类器，对未知图片进行预测
    for image_file in os.listdir(test_dir):
        #os.path.join连接路径名
        full_file_path = os.path.join(test_dir, image_file)
        print("Looking for faces in {}".format(image_file))
        # 使用训练模型进行预测
        predictions = predict(full_file_path, model_path="trained_knn_model.clf")
        # 输出结果
        return_name = []
        for name, (top, right, bottom, left) in predictions:
            print(type(name))
            print(111111111111)
            return_name.append(name)
            print("- Found {} at ({}, {})".format(name, left, top))
        # 图片上显示输出结果
        show_prediction_labels_on_image(os.path.join(test_dir, image_file), predictions,img_path_test,img_path_res)
    return return_name

